package com.okex.demo.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.event.EventListener;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.messaging.SessionConnectedEvent;
import org.springframework.web.socket.messaging.SessionDisconnectEvent;
import org.springframework.web.socket.messaging.SessionSubscribeEvent;
import org.springframework.web.socket.messaging.SessionUnsubscribeEvent;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * WebSocket事件监听器
 * 监听WebSocket的连接和断开事件，以便跟踪活跃用户和订阅
 * 使用事件机制通知数据服务启动和停止
 */
@Slf4j
@Component
public class WebSocketEventListener {

    private final ApplicationEventPublisher eventPublisher;

    // 记录活跃连接数
    private final AtomicInteger activeConnections = new AtomicInteger(0);
    
    // 记录每个产品的订阅数
    private final Map<String, AtomicInteger> productSubscriptions = new ConcurrentHashMap<>();
    
    // 记录每个会话ID和产品的关系
    private final Map<String, String> sessionSubscriptions = new ConcurrentHashMap<>();
    
    public WebSocketEventListener(ApplicationEventPublisher eventPublisher) {
        this.eventPublisher = eventPublisher;
    }
    
    /**
     * 处理连接事件
     */
    @EventListener
    public void handleWebSocketConnectListener(SessionConnectedEvent event) {
        StompHeaderAccessor headerAccessor = StompHeaderAccessor.wrap(event.getMessage());
        String sessionId = headerAccessor.getSessionId();
        
        int currentConnections = activeConnections.incrementAndGet();
        log.info("WebSocket连接建立: 会话ID {}, 当前连接数 {}", sessionId, currentConnections);
        
        // 首次连接时，发布连接事件
        if (currentConnections == 1) {
            log.info("首个客户端连接，发布连接事件");
            eventPublisher.publishEvent(new WebSocketConnectionEvent(this, true));
        }
    }
    
    /**
     * 处理断开事件
     */
    @EventListener
    public void handleWebSocketDisconnectListener(SessionDisconnectEvent event) {
        StompHeaderAccessor headerAccessor = StompHeaderAccessor.wrap(event.getMessage());
        String sessionId = headerAccessor.getSessionId();
        
        // 移除该会话的所有订阅
        removeSessionSubscriptions(sessionId);
        
        int currentConnections = activeConnections.decrementAndGet();
        log.info("WebSocket连接断开: 会话ID {}, 当前连接数 {}", sessionId, currentConnections);
        
        // 最后一个连接断开时，发布断开事件
        if (currentConnections <= 0) {
            log.info("所有客户端已断开，发布断开事件");
            eventPublisher.publishEvent(new WebSocketConnectionEvent(this, false));
            
            // 确保连接计数不会为负数
            if (currentConnections < 0) {
                activeConnections.set(0);
            }
        }
    }
    
    /**
     * 处理订阅事件
     */
    @EventListener
    public void handleWebSocketSubscribeListener(SessionSubscribeEvent event) {
        StompHeaderAccessor headerAccessor = StompHeaderAccessor.wrap(event.getMessage());
        String destination = headerAccessor.getDestination();
        String sessionId = headerAccessor.getSessionId();
        
        if (destination != null && destination.startsWith("/topic/ticker/")) {
            // 提取产品ID
            String productId = destination.substring("/topic/ticker/".length());
            incrementProductSubscription(productId);
            
            // 记录会话订阅
            sessionSubscriptions.put(sessionId + "_" + productId, productId);
            
            log.info("新增产品订阅: 会话ID {}, 产品 {}, 当前订阅数: {}", 
                    sessionId, productId, getProductSubscriptionCount(productId));
            
            // 发布产品订阅事件
            eventPublisher.publishEvent(new ProductSubscriptionEvent(this, productId, true));
        }
    }
    
    /**
     * 处理取消订阅事件
     */
    @EventListener
    public void handleWebSocketUnsubscribeListener(SessionUnsubscribeEvent event) {
        StompHeaderAccessor headerAccessor = StompHeaderAccessor.wrap(event.getMessage());
        String destination = headerAccessor.getDestination();
        String sessionId = headerAccessor.getSessionId();
        
        if (destination != null && destination.startsWith("/topic/ticker/")) {
            // 提取产品ID
            String productId = destination.substring("/topic/ticker/".length());
            decrementProductSubscription(productId);
            
            // 移除会话订阅记录
            sessionSubscriptions.remove(sessionId + "_" + productId);
            
            log.info("取消产品订阅: 会话ID {}, 产品 {}, 当前订阅数: {}", 
                    sessionId, productId, getProductSubscriptionCount(productId));
            
            // 发布产品取消订阅事件
            eventPublisher.publishEvent(new ProductSubscriptionEvent(this, productId, false));
        }
    }
    
    /**
     * 移除会话的所有订阅
     */
    private void removeSessionSubscriptions(String sessionId) {
        if (sessionId == null) {
            return;
        }
        
        // 找出该会话的所有订阅 (复制一份避免ConcurrentModificationException)
        Map<String, String> toRemove = new ConcurrentHashMap<>();
        
        sessionSubscriptions.entrySet().stream()
            .filter(entry -> entry.getKey().startsWith(sessionId + "_"))
            .forEach(entry -> toRemove.put(entry.getKey(), entry.getValue()));
        
        // 处理每一个需要移除的订阅
        for (Map.Entry<String, String> entry : toRemove.entrySet()) {
            String productId = entry.getValue();
            decrementProductSubscription(productId);
            sessionSubscriptions.remove(entry.getKey());
            log.info("断开连接移除订阅: 会话ID {}, 产品 {}, 当前订阅数: {}", 
                    sessionId, productId, getProductSubscriptionCount(productId));
        }
    }
    
    /**
     * 增加产品订阅计数
     */
    private void incrementProductSubscription(String productId) {
        productSubscriptions.computeIfAbsent(productId, k -> new AtomicInteger(0)).incrementAndGet();
    }
    
    /**
     * 减少产品订阅计数
     */
    private void decrementProductSubscription(String productId) {
        AtomicInteger count = productSubscriptions.get(productId);
        if (count != null) {
            int newCount = count.decrementAndGet();
            if (newCount <= 0) {
                productSubscriptions.remove(productId);
            }
        }
    }
    
    /**
     * 获取总连接数
     */
    public int getActiveConnectionCount() {
        return activeConnections.get();
    }
    
    /**
     * 检查是否有活跃连接
     */
    public boolean hasActiveConnections() {
        return activeConnections.get() > 0;
    }
    
    /**
     * 获取某个产品的订阅数
     */
    public int getProductSubscriptionCount(String productId) {
        AtomicInteger count = productSubscriptions.get(productId);
        return count != null ? count.get() : 0;
    }
    
    /**
     * 检查产品是否有订阅者
     */
    public boolean hasProductSubscriptions(String productId) {
        return getProductSubscriptionCount(productId) > 0;
    }
    
    /**
     * 获取所有订阅的产品ID
     */
    public Iterable<String> getSubscribedProducts() {
        return productSubscriptions.keySet();
    }
    
    /**
     * WebSocket连接事件，用于通知外部服务
     */
    public static class WebSocketConnectionEvent {
        private final Object source;
        private final boolean connected;
        
        public WebSocketConnectionEvent(Object source, boolean connected) {
            this.source = source;
            this.connected = connected;
        }
        
        public Object getSource() {
            return source;
        }
        
        public boolean isConnected() {
            return connected;
        }
    }
    
    /**
     * 产品订阅事件，用于通知外部服务
     */
    public static class ProductSubscriptionEvent {
        private final Object source;
        private final String productId;
        private final boolean subscribed;
        
        public ProductSubscriptionEvent(Object source, String productId, boolean subscribed) {
            this.source = source;
            this.productId = productId;
            this.subscribed = subscribed;
        }
        
        public Object getSource() {
            return source;
        }
        
        public String getProductId() {
            return productId;
        }
        
        public boolean isSubscribed() {
            return subscribed;
        }
    }
} 